import time
import numpy as np
import torch
import torch.nn as nn
from torch.utils import tensorboard
from torch.utils.data import DataLoader
import torch.optim as optim
import pickle
import copy
import os.path as osp
import os
from utils.utility import (display_performance, 
                           getLogger,
                           parallelize_model)
from models.convenience import (TrainMode,
                                ObsEncoderTrainer)



logger = getLogger(__name__)


# TODO: use learning rate sceduler
def train_encoder(net, train_loader, val_loader, perf_dict, save_dir_base, prefix, loss_fn, args, device):
    batch_num = len(train_loader)
    optimizer = optim.Adam(net.parameters(), lr=args.base_lar)
    dict_key = prefix + "_Encoder"

    start_time = time.time()
    # The train steps
    # TODO: define the train steps
    for epoch in range(args.epochs):
        train_loss, val_loss = 0, 0
        # TODO: compute metrics
        # val_res = np.zeros(2).astype(float)
        net.train()
        for data in train_loader:
            batch_data, batch_data_gt = data["observation"].to(device), data["prediction"].to(device)
            optimizer.zero_grad()
            net_output = net(batch_data)
            loss = loss_fn(net_output, batch_data_gt)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
        train_loss /= batch_num

        net.eval()
        with torch.no_grad():
            for data in val_loader:
                batch_data, batch_data_gt = data["observation"].to(device), data["prediction"].to(device)
                net_output = net(batch_data)
                loss = loss_fn(net_output, batch_data_gt)
                val_loss += loss.item()
                # TODO: compute metrics
                



def train(train_set, val_set, op_code, dataset_name, args, save_model_dir, device):

    # The performance metrics file
    perf_file_path = osp.join(save_model_dir, "Performances.pkl")
    if osp.isdir(perf_file_path):
        perf_dict = None
        with open(perf_file_path, "rb") as perf_f:
            perf_dict = pickle.load(perf_f)
        display_performance(perf_dict)
    else:
        perf_dict = {
            "Obs_Encoder": [1e3, 1e3],
            "Pred_Encoder": [1e3, 1e3],
            "Decoder": [[1e3, 1e3]] * 3,
            "Synthesizer": [[1e3, 1e3]] * 4
        }

    if op_code == TrainMode.TRAIN_OBS_ENC:
        if args.retrain:
            logger.info("Retrain Obs Encoder")
            perf_dict["Obs_Encoder"] = [1e3, 1e3]

        logger.info("Training Obs_Encoder...")
        start_time = time.time()
        net = ObsEncoderTrainer(obs_len=args.obs_len,
                                pred_len=args.pred_len,
                                hidden_dim=args.obs_feat_channels)
        net.to(device)
        net = parallelize_model(net)
        net.to(device)
        train_set.set_mode(TrainMode.TRAIN_OBS_ENC)
        val_set.set_mode(TrainMode.TRAIN_OBS_ENC)
        train_loader = DataLoader(train_set, batch_size=args.batch_size)
        val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False)
        
        # TODO: define the train_encoder encoder trainer
        train_encoder(net, 
                      train_loader, 
                      val_loader, 
                      perf_dict, 
                      save_model_dir, 
                      "Obs", 
                      loss_fn = nn.MSELoss(), 
                      args=args, 
                      device=device)
        end_time = time.time()
        logger.info(f"Time: {end_time-start_time:.2f}")

        
